import paddle
import paddle.nn as nn
import numpy as np
from paddle.vision.transforms import Compose, Resize, ToTensor, Normalize
import cv2

import paddle
from paddle.io import Dataset
import os
from PIL import Image

# 定义 CustomDataset 类（需放在文件顶部或使用前）
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.classes = sorted(os.listdir(data_dir))
        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}
        self.samples = self._load_samples()

    def _load_samples(self):
        samples = []
        for cls in self.classes:
            cls_path = os.path.join(self.data_dir, cls)
            for img_name in os.listdir(cls_path):
                img_path = os.path.join(cls_path, img_name)
                label = self.class_to_idx[cls]
                samples.append((img_path, label))
        return samples

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, paddle.to_tensor(label, dtype='int64')

# 使用 CustomDataset
train_data_dir = './dataset/train'
train_transforms = paddle.vision.transforms.Compose([...])
train_dataset = CustomDataset(data_dir=train_data_dir, transform=train_transforms)
# 定义 TNT 模型
class TNT(nn.Layer):
    def __init__(self, num_classes=10):
        super(TNT, self).__init__()
        # 这里假设需要将图像展平为一维向量，然后输入到全连接层
        self.flatten = nn.Flatten()
        # 计算展平后的维度，224*224*3 是因为图像大小为 224x224，通道数为 3
        flattened_dim = 224 * 224 * 3
        self.fc = nn.Linear(flattened_dim, num_classes)

    def forward(self, x):
        # 先将图像数据展平
        x = self.flatten(x)
        # 再通过全连接层
        x = self.fc(x)
        return x


# 创建模型实例
model = TNT(num_classes=2)

# 加载模型参数
params = paddle.load('./save_models/final.pdparams')
model.set_state_dict(params)
model.eval()  # 设置模型为评估模式

# 定义数据预处理管道
transform = Compose([
    Resize((224, 224)),  # 调整图像大小
    ToTensor(),  # 转换为张量
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 读取图像
image = cv2.imread(r'D:\TNT\dataset\test\cat\cat.109.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 预处理图像
input_image = transform(image).unsqueeze(0)  # 添加批次维度

# 进行预测
with paddle.no_grad():
    output = model(input_image)

# 获取预测结果
predicted_class = paddle.argmax(output, axis=1).item()
print(f"预测的类别是: {predicted_class}")
dataset = CustomDataset(data_dir='./dataset/train')
print("类别名称:", dataset.classes)       # 输出: ['cat', 'dog'] 或 ['dog', 'cat']
print("标签映射:", dataset.class_to_idx)  # 输出: {'cat': 0, 'dog': 1} 或 {'dog': 0, 'cat': 1}
